// Copyright (C) Kumo inc. and its affiliates.
// Author: Jeff.li lijippy@163.com
// All rights reserved.
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU Affero General Public License as published
// by the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public License
// along with this program.  If not, see <https://www.gnu.org/licenses/>.
//


#pragma once

#include <pollux/dwio/common/selective_struct_column_reader.h>
#include <pollux/dwio/parquet/common/level_conversion.h>

namespace kumo::pollux::dwio::common {
    class BufferedInput;
}

namespace kumo::pollux::parquet {
    enum class LevelMode;
    class PageReader;
    class ParquetParams;

    class StructColumnReader : public dwio::common::SelectiveStructColumnReader {
    public:
        StructColumnReader(
            const TypePtr &requestedType,
            const std::shared_ptr<const dwio::common::TypeWithId> &fileType,
            ParquetParams &params,
            common::ScanSpec &scanSpec);

        void read(int64_t offset, const RowSet &rows, const uint64_t *incomingNulls)
        override;

        void seekToRowGroup(int64_t index) override;

        /// Creates the streams for 'rowGroup'. Checks whether row 'rowGroup'
        /// has been buffered in 'input'. If true, return the input. Or else creates
        /// the streams in a new input and loads.
        std::shared_ptr<dwio::common::BufferedInput> loadRowGroup(
            uint32_t index,
            const std::shared_ptr<dwio::common::BufferedInput> &input);

        // No-op in Parquet. All readers switch row groups at the same time, there is
        // no on-demand skipping to a new row group.
        void advanceFieldReader(
            dwio::common::SelectiveColumnReader * /*reader*/,
            int64_t /*offset*/) override {
        }

        void setNullsFromRepDefs(PageReader &pageReader);

        dwio::common::SelectiveColumnReader *childForRepDefs() const {
            return childForRepDefs_;
        }

        /// Nested struct readers all get null flags and lengths for
        /// contained repeated readers for each range of top level rows. At
        /// the end of a read() with filters in different members, some of
        /// wich are structs themselves, different inner structs may be left
        /// on different rows. Before receiving the next set of
        /// nulls/lengths, the contained complex readers need to be
        /// positioned at the end of the last set of nulls/lengths.
        void seekToEndOfPresetNulls();

        void filterRowGroups(
            uint64_t rowGroupSize,
            const dwio::common::StatsContext &,
            dwio::common::FormatData::FilterRowGroupsResult &) const override;

    private:
        dwio::common::SelectiveColumnReader *findBestLeaf();

        void enqueueRowGroup(uint32_t index, dwio::common::BufferedInput &input);

        bool isRowGroupBuffered(uint32_t index, dwio::common::BufferedInput &input);

        // Leaf column reader used for getting nullability information for
        // 'this'. This is nullptr for the root of a table.
        dwio::common::SelectiveColumnReader *childForRepDefs_{nullptr};

        // Mode for getting nulls from repdefs. kStructOverLists if 'this'
        // only has list children.
        LevelMode levelMode_;

        // The level information for extracting nulls for 'this' from the
        // repdefs in a leaf PageReader.
        LevelInfo levelInfo_;
    };
} // namespace kumo::pollux::parquet
